/*
    Copyright 2007-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Implementation and explicit instantiations of the
/// ir_intensity_stage class. \ingroup mcrx

// $Id$

#include "config.h"
#include "mcrx-stage.h"
#include "mcrx-stage-impl.h"
#include "grid-fits.h"
#include "create_grid.h"
#include "equilibrium.h"
#include "arepo_grid.h"
#include "arepo_emission.h"

using namespace blitz;

template <template <typename> class grid_type> 
bool mcrx::ir_intensity_stage<grid_type>::check_stage_state () const
{
  bool state = false;
  try {
    this->output_file_->pHDU().readKey("MCRX" + stage_ID(), state);
  }
  catch (...) {}
  return state;
}

template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::set_stage_state (bool state)
{
  this->output_file_->pHDU().addKey("MCRX" + stage_ID(), state, "");
}

template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::basic_setup ()
{
  if(this->n_rays_>0)
    cout << "Increasing the number of rays in the equilibrium calculation is not possible.\nCalculation will be restarted with the larger number of rays.\n";

  // read original and intensity wavelength vectors
  CCfits::ExtHDU& lambda_hdu = open_HDU (*this->output_file_, "LAMBDA");
  array_1 lambda_orig;
  read(lambda_hdu.column("lambda"), lambda_orig);
  read(lambda_hdu.column("lambda_intensity"), lambda_);
  int n_lambda_int;
  lambda_hdu.readKey("NLAMBDA_INTENSITY", n_lambda_int);
  lambda_.resizeAndPreserve(n_lambda_int);

  // Read dust grains
  cout << "Setting up dust model" << endl;
  typename T_dust_model::T_scatterer_vector 
    sv(read_dust_grains<typename T_dust_model::T_scatterer_vector>(this->p_, this->m_.units));
  this->model_.reset(new T_dust_model (sv));
  this->model_->set_wavelength(lambda_);

  // create shadow vector of grain_model* for temp calculation
  for(int i=0; i< this->model_->n_scatterers(); ++i) {
    grain_models_.push_back
      (dynamic_cast<grain_model<polychromatic_scatterer_policy, 
       mcrx_rng_policy>*>(&this->model_->get_scatterer(i)));
    grain_models_.back()->resample(lambda_);
    // trigger generation of interpolation table
    //grain_models_.back()->calculate_SED(lambda_,0);
  }

  // Now build grid
  cout << "Building grid" << endl;
  
  // The density generator is used to generate the densities of the
  // dust grains from the gas and metal densities in the simulation.  
  uniform_density_generator dg(this->p_);

  T_grid* dummy=0;
  this->g_ = create_grid(*this->input_file_, dg, lambda_.size(), dummy);

  // Create ir_grid emission object which calculates the dust emission
  // in the grid cells (temp calc is not done here)
  cout << "Setting up emission" << endl;

  T_emission* dummy2=0;
  this->emi_ = create_grid(*this->input_file_, dg, lambda_.size(), dummy2);

  // This stage uses an empty emergence, as we are only calculating intensities
  this->eme_.reset(new T_emergence (vector<std::pair<T_float, T_float> >(),
			      1,1,1));

  // set reference wavelength fixed
  vector<T_float> templambda(lambda_.begin(), lambda_.end());
  reflambda_= 
    lower_bound (templambda.begin(), templambda.end(),
		 this->p_.getValue("ir_reference_wavelength", T_float())) -
    templambda.begin() - 1;
  cout << "Reference wavelength is " << lambda_(reflambda_) << " ("
       << reflambda_ << ")" << endl;

  CCfits::ExtHDU& intensity_hdu = open_HDU(*this->output_file_, "INTENSITY");
  read(intensity_hdu, star_intensity_);
  cout << "Allocating a " << star_intensity_.shape() 
       << " memory block for stellar intensities," 
       << product(star_intensity_.shape())*sizeof(T_float)*1.0/(1024*1024*1024)
       << " GB.\n";

  star_intensity_= pow (10., star_intensity_);

  cout << "Allocating a " << star_intensity_.shape() 
       << " memory block for dust intensities," 
       << product(star_intensity_.shape())*sizeof(T_float)*1.0/(1024*1024*1024)
       << " GB.\n";
  dust_intensity_.resize(star_intensity_.shape());
}

template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::setup_objects ()
{
  basic_setup();
}



/** Load IR stage data from a previously completed file. This is
     actually not supported. Because the temperature calculation have
     to be redone to convergence, there is very little lost by
     rerunning all rays. */
template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::load_file ()
{
  basic_setup();
}

/** Load optically thick IR stage data from a dump file.  */ 
template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::load_dump (binifstream& dump_file)
{
  basic_setup();

  // CHECK THAT THIS WORKS
  if (!this->m_.terminator()) {
    // Load dust temp distribution and cells remaining in calculation
    this->emi_->load_dump(dump_file);
    // load current grid intensities
    this->g_->load_dump(dump_file);

    // load dust_intensity
    TinyVector<int, 2> extent;
    dump_file >> extent [0]
	      >> extent [1];
    assert(all( dust_intensity_.shape() == extent));
    assert(dust_intensity_.isStorageContiguous());
    dump_file.read(reinterpret_cast<char*> (dust_intensity_.dataFirst()),
		   dust_intensity_.size()*sizeof (array_2::T_numtype) );
  }
}


/** Save cell intensities.  */
template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::save_file ()
{
  // Intensity is already normalized in equilibrium calculation. Note
  // that we want to save the dust_intensity_ array, not the intensity
  // in the grid object. The latter only includes the last iteration.
  bool do_compress =this->p_.defined("compress_images")&&
    this->p_.getValue("compress_images", bool ());

  const string i_name="DUSTINTENSITY";
  CCfits::ExtHDU* i_hdu;    
  try {
    i_hdu = &open_HDU (*this->output_file_, i_name);
  }
  catch (CCfits::FITS::NoSuchHDU&) {
    std::vector<long> naxes,znaxes;
    naxes.push_back(dust_intensity_.extent(firstDim));
    naxes.push_back(dust_intensity_.extent(secondDim));

    if (do_compress) {
      // use the cfitsio tile compression on these images
      znaxes=naxes;
      znaxes[0]=100; // compress 100 cells as one
      this->output_file_->setCompressionType(GZIP_1);
      this->output_file_->setTileDimensions(znaxes);
    }
    
    i_hdu = this->output_file_->addImage (i_name, FLOAT_IMG, naxes);
    this->output_file_->setCompressionType(0);
    
    i_hdu->writeComment("This HDU contains the mean intensities in the grid cells. The first dimension is grid cell, the second wavelength.");
  }
  std::cout << "Writing intensities" << std::endl;
  assert(dust_intensity_.size()>0);

  Array<float, 2> fintensity(dust_intensity_.shape());

  // handle zero-intensity case
  fintensity=where(dust_intensity_>0, 
		   log10(dust_intensity_), 
		   blitz::neghuge(float()));
  write(*i_hdu, fintensity);
}


/** Dump IR stage info. We might have been interrupted during the temp
    calculation, or we might have been interrupted during shooting. In
    either case, we should save the cell SEDs that have been
    calculated so we don't have to do it again. */
template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::save_dump (binofstream& dump_file) const
{
  this->g_->write_dump(dump_file);
}


template <template <typename> class grid_type> 
template<typename T>
void mcrx::ir_intensity_stage<grid_type>::calculate_dust_SED(const blitz::ETBase<T>& intensity)
{
  assert(0);
  mcrx_terminator term(this->m_);
  const int n_threads = this->p_.defined("n_threads")?
    this->p_.getValue("n_threads", int ()): 1;

  const bool check_energy = this->p_.defined("check_energy_conservation") ?
    this->p_.getValue("check_energy_conservation", bool()) : false;

  this->emi_->calculate_SED(intensity, grain_models_, term, n_threads, 0,
		      false);
}


template <template <typename> class grid_type> 
bool mcrx::ir_intensity_stage<grid_type>::shoot ()
{
  const T_float i_ls = this->p_.getValue("i_last_scatter", T_float ());
  const T_float i_s = this->p_.getValue("i_split", T_float ());
  const T_float ir_tol = 
    this->p_.getValue("ir_equilibrium_tolerance", T_float());
  const int n_f = this->p_.getValue("n_forced_ir", int(0), true);
  const int n_threads = this->p_.getValue("n_threads", int (1), true);
  assert (n_threads == this->m_.rng_states.size()); 
  const bool bind_threads = 
    this->p_.getValue("bind_threads", bool (false), true);
  const bool bank_size = this->p_.getValue("bank_size", int(1), true);

  mcrx_terminator term(this->m_);
  bool t;

  {
    // use weak references to avoid locking issues
    array_2 star_intensity_ref, dust_intensity_ref;
    star_intensity_ref.weakReference(star_intensity_);
    dust_intensity_ref.weakReference(dust_intensity_);
    
    t= determine_dust_equilibrium_intensities
      (
       *this->model_,
       *this->g_,
       *this->emi_,
       lambda_,
       star_intensity_ref,
       dust_intensity_ref,
       this->m_.rng_states,
       this->n_rays_desired_,
       n_threads,
       bind_threads,
       bank_size,
       ir_tol,
       i_ls,
       n_f,
       i_s,
       shooter(),
       term,
       false);
  }

  // set number of rays manually so it gets written out correctly
  this->n_rays_ = this->n_rays_desired_;
  // we can free the star intensity now
  star_intensity_.free();

  return t;
}

template <template <typename> class grid_type> 
void mcrx::ir_intensity_stage<grid_type>::operator() ()
{
  // this is here so the function is instantiated
  this->run_stage();
}


// explicit instantiations
template class mcrx::ir_intensity_stage<mcrx::adaptive_grid>;

#ifdef WITH_AREPO
template class mcrx::ir_intensity_stage<mcrx::arepo_grid>;
#endif
